#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
独立的传感器预测模型训练脚本
完全独立，不依赖其他模块，一键训练传感器预测模型
"""

import pandas as pd
import numpy as np
import os
import pickle
import warnings
from datetime import datetime, timedelta
from sklearn.preprocessing import MinMaxScaler
import tensorflow as tf

warnings.filterwarnings('ignore')
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

def create_sensor_model(input_shape):
    """创建传感器预测LSTM模型"""
    model = tf.keras.Sequential([
        tf.keras.layers.LSTM(128, return_sequences=True, input_shape=input_shape),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.LSTM(64, return_sequences=True),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.LSTM(32),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(96 * 21, activation='linear'),  # 96个时间点 × 21个特征
        tf.keras.layers.Reshape((96, 21))
    ])
    
    model.compile(
        optimizer='adam',
        loss='mse',
        metrics=['mae']
    )
    
    return model

def prepare_time_features(datetime_series):
    """准备时间特征"""
    time_features = pd.DataFrame()
    time_features['Hour_sin'] = np.sin(2 * np.pi * datetime_series.dt.hour / 24)
    time_features['Hour_cos'] = np.cos(2 * np.pi * datetime_series.dt.hour / 24)
    time_features['DayOfWeek_sin'] = np.sin(2 * np.pi * datetime_series.dt.dayofweek / 7)
    time_features['DayOfWeek_cos'] = np.cos(2 * np.pi * datetime_series.dt.dayofweek / 7)
    time_features['IsWeekend'] = (datetime_series.dt.dayofweek >= 5).astype(int)
    
    return time_features.values

def prepare_training_data(data_path, sequence_length=300, prediction_horizon=96):
    """准备训练数据"""
    print("📊 加载和预处理数据...")
    
    # 读取数据
    df = pd.read_csv(data_path)
    df['DateTime'] = pd.to_datetime(df['Date'] + ' ' + df['Time'], format='%d-%m-%Y %H:%M:%S')
    df = df.sort_values('DateTime').reset_index(drop=True)
    
    print(f"📅 数据时间范围: {df['DateTime'].min()} 到 {df['DateTime'].max()}")
    print(f"📈 数据总数: {len(df)} 条记录")
    
    # 传感器列
    sensor_columns = [
        'S1_Temp', 'S2_Temp', 'S3_Temp', 'S4_Temp',
        'S1_Light', 'S2_Light', 'S3_Light', 'S4_Light',
        'S1_Sound', 'S2_Sound', 'S3_Sound', 'S4_Sound',
        'S5_CO2', 'S5_CO2_Slope', 'S6_PIR', 'S7_PIR'
    ]
    
    # 提取传感器数据
    sensor_data = df[sensor_columns].values.astype(np.float32)
    
    # 准备时间特征
    time_features = prepare_time_features(df['DateTime'])
    
    # 合并特征
    combined_features = np.hstack([sensor_data, time_features])
    print(f"🔧 合并特征形状: {combined_features.shape}")
    
    # 数据标准化
    scaler = MinMaxScaler()
    normalized_data = scaler.fit_transform(combined_features)
    
    # 创建训练序列
    X, y = [], []
    
    for i in range(len(normalized_data) - sequence_length - prediction_horizon + 1):
        # 输入序列：过去sequence_length个时间点
        X.append(normalized_data[i:i + sequence_length])
        
        # 输出序列：未来prediction_horizon个时间点
        y.append(normalized_data[i + sequence_length:i + sequence_length + prediction_horizon])
    
    X = np.array(X)
    y = np.array(y)
    
    print(f"🎯 训练数据形状: X={X.shape}, y={y.shape}")
    
    return X, y, scaler, sensor_columns

def train_model():
    """训练传感器预测模型"""
    print("🚀 开始训练传感器预测模型...")
    
    # 检查数据文件
    data_path = 'data/room_occupancy.csv'
    if not os.path.exists(data_path):
        print(f"❌ 数据文件不存在: {data_path}")
        return
    
    # 准备训练数据
    X, y, scaler, sensor_columns = prepare_training_data(data_path)
    
    # 划分训练集和验证集
    split_idx = int(len(X) * 0.8)
    X_train, X_val = X[:split_idx], X[split_idx:]
    y_train, y_val = y[:split_idx], y[split_idx:]
    
    print(f"📊 训练集: {X_train.shape}, 验证集: {X_val.shape}")
    
    # 创建模型
    model = create_sensor_model((X.shape[1], X.shape[2]))
    print("🏗️  模型架构:")
    model.summary()
    
    # 训练回调
    callbacks = [
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss',
            patience=10,
            restore_best_weights=True
        ),
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss',
            factor=0.5,
            patience=5,
            min_lr=1e-6
        )
    ]
    
    # 训练模型
    print("🎯 开始训练...")
    history = model.fit(
        X_train, y_train,
        validation_data=(X_val, y_val),
        epochs=50,
        batch_size=16,
        callbacks=callbacks,
        verbose=1
    )
    
    # 保存模型和缩放器
    print("💾 保存模型...")
    
    # 确保保存目录存在
    save_dir = 'model/saved_models'
    os.makedirs(save_dir, exist_ok=True)
    
    # 保存模型
    model_path = os.path.join(save_dir, 'sensor_predictor.h5')
    model.save(model_path)
    print(f"✅ 模型已保存: {model_path}")
    
    # 保存缩放器
    scalers_path = os.path.join(save_dir, 'sensor_scalers.pkl')
    scalers_data = {
        'combined_scaler': scaler,
        'sensor_columns': sensor_columns
    }
    
    with open(scalers_path, 'wb') as f:
        pickle.dump(scalers_data, f)
    print(f"✅ 缩放器已保存: {scalers_path}")
    
    # 测试预测功能
    print("\n🧪 测试预测功能...")
    test_prediction(model, scaler, sensor_columns)
    
    print("\n🎉 训练完成！")
    return model, scaler

def test_prediction(model, scaler, sensor_columns):
    """测试预测功能"""
    try:
        # 读取测试数据
        df = pd.read_csv('data/room_occupancy.csv')
        df['DateTime'] = pd.to_datetime(df['Date'] + ' ' + df['Time'], format='%d-%m-%Y %H:%M:%S')
        df = df.sort_values('DateTime').tail(300).reset_index(drop=True)
        
        # 准备输入数据
        sensor_data = df[sensor_columns].values.astype(np.float32)
        time_features = prepare_time_features(df['DateTime'])
        combined_features = np.hstack([sensor_data, time_features])
        
        # 标准化
        normalized_data = scaler.transform(combined_features)
        input_sequence = normalized_data.reshape(1, 300, -1)
        
        # 预测
        predictions = model.predict(input_sequence, verbose=0)
        
        # 反标准化
        denormalized_predictions = scaler.inverse_transform(predictions[0])
        sensor_predictions = denormalized_predictions[:, :len(sensor_columns)]
        
        # 生成未来时间
        last_datetime = df['DateTime'].max()
        future_times = []
        current_time = last_datetime + timedelta(minutes=30)
        
        for i in range(96):
            future_times.append(current_time)
            current_time += timedelta(minutes=30)
        
        # 创建结果DataFrame
        results = []
        for i, future_time in enumerate(future_times):
            row = {
                'Date': future_time.strftime('%d-%m-%Y'),
                'Time': future_time.strftime('%H:%M:%S')
            }
            
            # 添加传感器预测值
            for j, col in enumerate(sensor_columns):
                row[col] = float(sensor_predictions[i, j])
            
            results.append(row)
        
        results_df = pd.DataFrame(results)
        
        # 应用约束
        results_df = apply_sensor_constraints(results_df)
        
        # 保存测试结果
        results_df.to_csv('predicted_sensors_test.csv', index=False)
        print(f"✅ 测试完成，结果已保存到 predicted_sensors_test.csv")
        
        # 显示示例
        print("📊 预测结果前5行:")
        print(results_df[['Date', 'Time', 'S1_Temp', 'S5_CO2']].head())
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")

def apply_sensor_constraints(df):
    """应用传感器数据约束"""
    # 温度传感器：15-35°C
    temp_cols = [col for col in df.columns if 'Temp' in col]
    for col in temp_cols:
        df[col] = np.clip(df[col], 15, 35)
    
    # 光线传感器：0-1000
    light_cols = [col for col in df.columns if 'Light' in col]
    for col in light_cols:
        df[col] = np.clip(df[col], 0, 1000)
    
    # 声音传感器：0-1
    sound_cols = [col for col in df.columns if 'Sound' in col]
    for col in sound_cols:
        df[col] = np.clip(df[col], 0, 1)
    
    # CO2传感器：300-1000 ppm
    if 'S5_CO2' in df.columns:
        df['S5_CO2'] = np.clip(df['S5_CO2'], 300, 1000)
    
    # CO2斜率：-50到50
    if 'S5_CO2_Slope' in df.columns:
        df['S5_CO2_Slope'] = np.clip(df['S5_CO2_Slope'], -50, 50)
    
    # PIR传感器：0或1
    pir_cols = [col for col in df.columns if 'PIR' in col]
    for col in pir_cols:
        df[col] = np.round(np.clip(df[col], 0, 1)).astype(int)
    
    return df

def main():
    """主函数"""
    print("=" * 60)
    print("🎯 独立传感器预测模型训练器")
    print("=" * 60)
    
    try:
        model, scaler = train_model()
        print("\n🎉 所有步骤完成！")
        print("📁 生成的文件:")
        print("   - model/saved_models/sensor_predictor.h5")
        print("   - model/saved_models/sensor_scalers.pkl")
        print("   - predicted_sensors_test.csv")
        
    except Exception as e:
        print(f"❌ 训练失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 